"""
Poisson Matrix Assembly

Assembles discrete Laplacian operators for the Poisson equation ∇²ψ = rhs
on L-shaped domains with appropriate boundary conditions. Supports both 
2D and 3D formulations with integration to GCR.jl solvers.
"""

"""
    assemble_poisson_matrix(domain::LShapeDomain{2}, h::Float64) -> SparseMatrixCSC{Float64, Int}

Assemble the 2D discrete Laplacian matrix for interior nodes in the L-shaped domain.

The matrix implements the 5-point finite difference stencil:
∇²ψ ≈ (ψ_{i+1,j} + ψ_{i-1,j} + ψ_{i,j+1} + ψ_{i,j-1} - 4ψ_{i,j})/h²

Boundary conditions:
- Dirichlet BC: ψ = 0 on fluid boundaries
- Mirror condition: ∂ψ/∂n = 0 on solid interfaces (effectively absorbed into interior stencil)

# Arguments
- `domain::LShapeDomain{2}`: L-shaped domain geometry
- `h::Float64`: Grid spacing

# Returns
- `A::SparseMatrixCSC{Float64, Int}`: Assembled Laplacian matrix (size n_interior × n_interior)
- Matrix is symmetric negative definite (eigenvalues < 0)
"""
function assemble_poisson_matrix(domain::LShapeDomain{2}, h::Float64)
    # Create mapping from interior nodes to compact matrix indices
    interior_indices = findall(domain.interior_mask)
    n_interior = length(interior_indices)
    
    # Create index mapping array for fast lookup
    id_map = zeros(Int, size(domain.interior_mask))
    for (k, idx) in enumerate(interior_indices)
        id_map[idx] = k
    end
    
    # Preallocate matrix assembly arrays with better size estimates
    # For 2D: each interior node has at most 5 connections (center + 4 neighbors)
    # Account for boundary modifications that may reduce connections
    max_entries = n_interior * 5
    I = Vector{Int}(undef, max_entries)
    J = Vector{Int}(undef, max_entries) 
    V = Vector{Float64}(undef, max_entries)
    
    entry_count = 0  # Track actual entries added
    
    h2 = h * h
    nx, ny = size(domain.fluid_mask)
    
    # Assemble matrix row by row
    for (row_id, idx) in enumerate(interior_indices)
        i, j = Tuple(idx)  # Convert linear index to (i,j) subscripts
        
        # Diagonal coefficient (starts at -4 for 5-point stencil)
        diag_coeff = -4.0
        
        # Check four neighbors: right, left, up, down
        neighbors = [(i+1, j), (i-1, j), (i, j+1), (i, j-1)]
        
        for (ni, nj) in neighbors
            # Skip if outside domain bounds
            if ni < 1 || ni > nx || nj < 1 || nj > ny
                continue
            end
            
            if domain.fluid_mask[ni, nj]
                if domain.interior_mask[ni, nj]
                    # Interior neighbor: add off-diagonal entry
                    neighbor_id = id_map[ni, nj]
                    entry_count += 1
                    I[entry_count] = row_id
                    J[entry_count] = neighbor_id 
                    V[entry_count] = 1.0
                else
                    # Boundary fluid node: Dirichlet BC (ψ = 0)
                    # No contribution to matrix (implicitly ψ = 0)
                end
            else
                # Solid neighbor: mirror condition ∂ψ/∂n = 0
                # This modifies the diagonal coefficient
                diag_coeff += 1.0
            end
        end
        
        # Add diagonal entry
        entry_count += 1
        I[entry_count] = row_id
        J[entry_count] = row_id
        V[entry_count] = diag_coeff / h2  # Scale by 1/h²
    end
    
    # Resize arrays to actual size and assemble sparse matrix
    resize!(I, entry_count)
    resize!(J, entry_count)
    resize!(V, entry_count)
    
    # Assemble sparse matrix with memory-efficient constructor
    A = sparse(I, J, V, n_interior, n_interior)
    
    # The assembled matrix represents -∇² (negative Laplacian)
    # For the Poisson equation ∇²ψ = rhs, we solve -A*ψ = -rhs
    # Since A has negative eigenvalues, -A is positive definite
    
    return A, id_map, interior_indices
end

"""
    assemble_poisson_matrix(domain::LShapeDomain{3}, h::Float64) -> SparseMatrixCSC{Float64, Int}

Assemble the 3D discrete Laplacian matrix using 7-point stencil.

The matrix implements:
∇²ψ ≈ (ψ_{i±1,j,k} + ψ_{i,j±1,k} + ψ_{i,j,k±1} - 6ψ_{i,j,k})/h²

# Arguments  
- `domain::LShapeDomain{3}`: 3D L-shaped domain
- `h::Float64`: Grid spacing

# Returns
- `A::SparseMatrixCSC{Float64, Int}`: 3D Laplacian matrix
"""
function assemble_poisson_matrix(domain::LShapeDomain{3}, h::Float64)
    interior_indices = findall(domain.interior_mask)
    n_interior = length(interior_indices)
    
    # Create 3D index mapping
    id_map = zeros(Int, size(domain.interior_mask))
    for (k, idx) in enumerate(interior_indices)
        id_map[idx] = k
    end
    
    # Preallocate arrays for 3D: each interior node has at most 7 connections
    max_entries = n_interior * 7
    I = Vector{Int}(undef, max_entries)
    J = Vector{Int}(undef, max_entries)
    V = Vector{Float64}(undef, max_entries)
    
    entry_count = 0
    
    h2 = h * h
    nx, ny, nz = size(domain.fluid_mask)
    
    for (row_id, idx) in enumerate(interior_indices)
        i, j, k = Tuple(idx)  # 3D subscripts
        
        # Diagonal coefficient (starts at -6 for 7-point stencil)
        diag_coeff = -6.0
        
        # Check six neighbors: ±x, ±y, ±z directions
        neighbors = [(i+1,j,k), (i-1,j,k), (i,j+1,k), (i,j-1,k), (i,j,k+1), (i,j,k-1)]
        
        for (ni, nj, nk) in neighbors
            # Bounds check
            if ni < 1 || ni > nx || nj < 1 || nj > ny || nk < 1 || nk > nz
                continue
            end
            
            if domain.fluid_mask[ni, nj, nk]
                if domain.interior_mask[ni, nj, nk]
                    # Interior neighbor
                    neighbor_id = id_map[ni, nj, nk]
                    entry_count += 1
                    I[entry_count] = row_id
                    J[entry_count] = neighbor_id
                    V[entry_count] = 1.0
                else
                    # Boundary fluid node: Dirichlet BC
                end
            else
                # Solid neighbor: mirror condition
                diag_coeff += 1.0
            end
        end
        
        # Add diagonal entry
        entry_count += 1
        I[entry_count] = row_id
        J[entry_count] = row_id
        V[entry_count] = diag_coeff / h2
    end
    
    # Resize arrays to actual size and assemble sparse matrix
    resize!(I, entry_count)
    resize!(J, entry_count) 
    resize!(V, entry_count)
    
    A = sparse(I, J, V, n_interior, n_interior)
    return A, id_map, interior_indices
end

"""
    assemble_poisson_matrix_efficient(domain::LShapeDomain{D}, h::Float64; 
                                     use_threads=true, batch_size=1000) where D

Memory-efficient assembly for large Poisson matrices using batch processing and optional threading.

This function is optimized for large-scale problems where memory efficiency is critical.
It uses batch processing to limit memory usage during assembly and can leverage
multiple threads for improved performance.

# Arguments
- `domain::LShapeDomain{D}`: Domain geometry
- `h::Float64`: Grid spacing
- `use_threads::Bool`: Whether to use multi-threading (default: true)
- `batch_size::Int`: Number of rows to process per batch (default: 1000)

# Returns
- Same as regular assemble_poisson_matrix but with optimized memory usage
"""
function assemble_poisson_matrix_efficient(domain::LShapeDomain{D}, h::Float64; 
                                          use_threads=true, batch_size=1000) where D
    
    interior_indices = findall(domain.interior_mask)
    n_interior = length(interior_indices)
    
    if n_interior <= 5000
        # For small problems, use regular assembly
        return assemble_poisson_matrix(domain, h)
    end
    
    @info "Using efficient assembly for large problem (n_interior = $n_interior)"
    
    # Create index mapping
    id_map = zeros(Int, size(domain.interior_mask))
    for (k, idx) in enumerate(interior_indices)
        id_map[idx] = k
    end
    
    # Estimate matrix sparsity for preallocation
    stencil_size = D == 2 ? 5 : 7
    estimated_nnz = min(n_interior * stencil_size, n_interior * (n_interior + 1) ÷ 2)
    
    # Use sparse matrix builder for efficient incremental construction
    builder = spzeros(Float64, n_interior, n_interior)
    
    # Process in batches to control memory usage
    h2 = h * h
    n_batches = ceil(Int, n_interior / batch_size)
    
    @info "Processing matrix assembly in $n_batches batches"
    
    for batch_idx in 1:n_batches
        start_row = (batch_idx - 1) * batch_size + 1
        end_row = min(batch_idx * batch_size, n_interior)
        
        batch_rows = start_row:end_row
        batch_I = Int[]
        batch_J = Int[] 
        batch_V = Float64[]
        
        for row_id in batch_rows
            idx = interior_indices[row_id]
            
            if D == 2
                i, j = Tuple(idx)
                diag_coeff = -4.0
                neighbors = [(i+1, j), (i-1, j), (i, j+1), (i, j-1)]
                domain_size = size(domain.fluid_mask)
            else  # 3D
                i, j, k = Tuple(idx)
                diag_coeff = -6.0
                neighbors = [(i+1,j,k), (i-1,j,k), (i,j+1,k), (i,j-1,k), (i,j,k+1), (i,j,k-1)]
                domain_size = size(domain.fluid_mask)
            end
            
            # Process neighbors
            for neighbor_coords in neighbors
                if all(1 ≤ neighbor_coords[d] ≤ domain_size[d] for d in 1:D)
                    if domain.fluid_mask[neighbor_coords...]
                        if domain.interior_mask[neighbor_coords...]
                            neighbor_id = id_map[neighbor_coords...]
                            push!(batch_I, row_id)
                            push!(batch_J, neighbor_id)
                            push!(batch_V, 1.0)
                        end
                    else
                        diag_coeff += 1.0
                    end
                end
            end
            
            # Add diagonal entry
            push!(batch_I, row_id)
            push!(batch_J, row_id) 
            push!(batch_V, diag_coeff / h2)
        end
        
        # Add batch entries to matrix
        batch_sparse = sparse(batch_I, batch_J, batch_V, n_interior, n_interior)
        builder += batch_sparse
        
        # Force garbage collection periodically for large problems
        if batch_idx % 10 == 0
            GC.gc()
        end
    end
    
    return builder, id_map, interior_indices
end

"""
    extract_interior_values(field::Array{Float64}, interior_indices::Vector{CartesianIndex}) -> Vector{Float64}

Extract values at interior nodes from a full field array.

# Arguments
- `field::Array{Float64}`: Full field array (2D or 3D)
- `interior_indices::Vector{CartesianIndex}`: Interior node indices

# Returns
- `values::Vector{Float64}`: Values at interior nodes (compact format)
"""
function extract_interior_values(field::Array{Float64}, interior_indices::Vector{CartesianIndex{N}}) where N
    return [field[idx] for idx in interior_indices]
end

"""
    scatter_interior_values!(field::Array{Float64}, values::Vector{Float64}, 
                            interior_indices::Vector{CartesianIndex})

Scatter values from compact format back to interior nodes of full field array.

# Arguments
- `field::Array{Float64}`: Full field array (modified in-place)
- `values::Vector{Float64}`: Compact values at interior nodes
- `interior_indices::Vector{CartesianIndex}`: Interior node indices
"""
function scatter_interior_values!(field::Array{Float64}, values::Vector{Float64}, 
                                 interior_indices::Vector{CartesianIndex{N}}) where N
    @assert length(values) == length(interior_indices)
    
    @inbounds for (k, idx) in enumerate(interior_indices)
        field[idx] = values[k]
    end
end

"""
    compute_matrix_properties(A::SparseMatrixCSC{Float64, Int}) -> NamedTuple

Compute diagnostic properties of the assembled Poisson matrix.

# Returns
Named tuple with matrix properties:
- `size`: Matrix dimensions
- `nnz`: Number of non-zeros
- `density`: Sparsity density 
- `is_symmetric`: Whether matrix is symmetric
- `min_eigenval_est`: Estimate of smallest eigenvalue
- `max_eigenval_est`: Estimate of largest eigenvalue
- `condition_est`: Condition number estimate
"""
function compute_matrix_properties(A::SparseMatrixCSC{Float64, Int})
    n = size(A, 1)
    nnz_A = nnz(A)
    density = nnz_A / n^2
    
    # Check symmetry (should be true for our Laplacian)
    is_symmetric = issymmetric(A)
    
    # For large matrices, estimate eigenvalues using power iteration or sampling
    # For small matrices, compute exactly
    if n <= 200
        try
            eigenvals = eigvals(Matrix(A))
            min_eval = minimum(real(eigenvals))
            max_eval = maximum(real(eigenvals))
            cond_est = -max_eval / min_eval  # Condition number for negative definite matrix
        catch
            # Fallback estimates
            min_eval = -8.0  # Rough estimate for 2D Laplacian  
            max_eval = -0.1
            cond_est = 80.0
        end
    else
        # Rough estimates for large matrices
        h_est = sqrt(4.0 / n)  # Estimate grid spacing from DOF count
        min_eval = -8.0 / h_est^2  # Largest magnitude eigenvalue
        max_eval = -π^2 / 4.0      # Smallest magnitude eigenvalue (rough)
        cond_est = min_eval / max_eval
    end
    
    return (
        size = (n, n),
        nnz = nnz_A,
        density = density,
        is_symmetric = is_symmetric,
        min_eigenval_est = min_eval,
        max_eigenval_est = max_eval, 
        condition_est = cond_est,
        memory_mb = (nnz_A * (8 + 4 + 4)) / (1024^2)  # Rough memory estimate
    )
end

"""
    validate_poisson_matrix(A::SparseMatrixCSC{Float64, Int}, domain::LShapeDomain{D}) -> Bool where D

Validate that the assembled Poisson matrix has expected properties.

# Arguments
- `A::SparseMatrixCSC{Float64, Int}`: Assembled matrix
- `domain::LShapeDomain{D}`: Domain geometry

# Returns  
- `is_valid::Bool`: Whether matrix passes validation checks
"""
function validate_poisson_matrix(A::SparseMatrixCSC{Float64, Int}, domain::LShapeDomain{D}) where D
    props = compute_matrix_properties(A)
    is_valid = true
    
    # Check basic properties
    expected_size = sum(domain.interior_mask)
    if props.size[1] != expected_size
        @error "Matrix size $(props.size) doesn't match interior nodes ($expected_size)"
        is_valid = false
    end
    
    if !props.is_symmetric
        @error "Poisson matrix should be symmetric"
        is_valid = false
    end
    
    # Check that matrix is negative definite (eigenvalues < 0)
    if props.max_eigenval_est >= 0
        @error "Matrix should be negative definite (max eigenvalue: $(props.max_eigenval_est))"
        is_valid = false
    end
    
    # Check sparsity pattern (should be reasonable for finite differences)
    expected_nnz_per_row = D == 2 ? 5 : 7  # Stencil size
    actual_nnz_per_row = props.nnz / props.size[1]
    
    if actual_nnz_per_row > 2 * expected_nnz_per_row
        @warn "Matrix seems too dense ($(actual_nnz_per_row) nnz/row, expected ~$expected_nnz_per_row)"
    end
    
    if props.condition_est > 1e10
        @warn "Matrix is poorly conditioned (condition ≈ $(props.condition_est))"
    end
    
    return is_valid
end

"""
    print_matrix_info(A::SparseMatrixCSC{Float64, Int}, domain::LShapeDomain{D}) where D

Print diagnostic information about the assembled Poisson matrix.
"""
function print_matrix_info(A::SparseMatrixCSC{Float64, Int}, domain::LShapeDomain{D}) where D
    props = compute_matrix_properties(A)
    
    println("Poisson Matrix Information ($(D)D):")
    println("  Size: $(props.size[1]) × $(props.size[2])")
    println("  Non-zeros: $(props.nnz) ($(round(100*props.density, digits=3))% density)")
    println("  Memory usage: $(round(props.memory_mb, digits=1)) MB")
    println("  Symmetric: $(props.is_symmetric)")
    println("  Eigenvalue range: [$(round(props.min_eigenval_est, digits=2)), $(round(props.max_eigenval_est, digits=2))]")
    println("  Condition number: $(round(props.condition_est, digits=1))")
    
    # Domain statistics
    total_nodes = length(domain.fluid_mask)
    interior_nodes = sum(domain.interior_mask)
    println("  Interior DOF: $interior_nodes / $total_nodes ($(round(100*interior_nodes/total_nodes, digits=1))%)")
    
    # Validation
    is_valid = validate_poisson_matrix(A, domain)
    println("  Matrix validation: $(is_valid ? "✓ PASSED" : "✗ FAILED")")
end